{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Concatenate all files\n",
    "\n",
    "```bash\n",
    "$ cd path/to/train-easy/\n",
    "$ find -name '*.txt' -exec cat {} \\; > ../../../interim/train-easy_all.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0-alpha0\n",
      "GPU Available:  True\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import glob\n",
    "import pickle\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "from lstm import LSTM_S2S\n",
    "from metrics import exact_match_metric\n",
    "from callbacks import NValidationSetsCallback, GradientLogger\n",
    "from generator import DataGenerator\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "print(tf.__version__)\n",
    "print(\"GPU Available: \", tf.test.is_gpu_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings_path = Path('../../settings/settings.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(str(settings_path), 'r') as file:\n",
    "    settings_dict = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'batch_size': 1024,\n",
       " 'epochs': 1,\n",
       " 'latent_dim': 256,\n",
       " 'math_module': 'arithmetic',\n",
       " 'save_path': '/artifacts/',\n",
       " 'thinking_steps': 1,\n",
       " 'train_level': 'easy'}"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "settings_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data\n",
    "\n",
    "Start with batching a single file before tackling the whole dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "extrapolate  interpolate  train-easy  train-hard  train-medium\n"
     ]
    }
   ],
   "source": [
    "raw_path = Path(settings_dict['data_path'])\n",
    "!ls {raw_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "algebra__linear_1d.txt\n",
      "algebra__linear_1d_composed.txt\n",
      "algebra__linear_2d.txt\n",
      "algebra__linear_2d_composed.txt\n",
      "algebra__polynomial_roots.txt\n"
     ]
    }
   ],
   "source": [
    "interpolate_path = raw_path/'interpolate'\n",
    "!ls {interpolate_path} | head -5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "algebra__polynomial_roots_big.txt\n",
      "arithmetic__add_or_sub_big.txt\n",
      "arithmetic__add_sub_multiple_longer.txt\n",
      "arithmetic__div_big.txt\n",
      "arithmetic__mixed_longer.txt\n"
     ]
    }
   ],
   "source": [
    "extrapolate_path = raw_path/'extrapolate'\n",
    "!ls {extrapolate_path} | head -5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "algebra__linear_1d.txt\n",
      "algebra__linear_1d_composed.txt\n",
      "algebra__linear_2d.txt\n",
      "algebra__linear_2d_composed.txt\n",
      "algebra__polynomial_roots.txt\n"
     ]
    }
   ],
   "source": [
    "train_easy_path = raw_path/'train-easy/'\n",
    "!ls {train_easy_path} | head -5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def concatenate_texts(path, pattern):\n",
    "    file_paths = list(path.glob('{}*.txt'.format(pattern)))\n",
    "    \n",
    "    input_texts = []\n",
    "    target_texts = []\n",
    "\n",
    "    for file_path in file_paths:\n",
    "        with open(str(file_path), 'r', encoding='utf-8') as f:\n",
    "            lines = f.read().split('\\n')[:-1]\n",
    "\n",
    "        input_texts.extend(lines[0::2])\n",
    "        target_texts.extend(['\\t' + target_text + '\\n' for target_text in lines[1::2]])\n",
    "        \n",
    "    return input_texts, target_texts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "math_module = settings_dict[\"math_module\"]\n",
    "train_level = settings_dict[\"train_level\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {\n",
    "    'train':(raw_path, 'train-' + train_level + '/' + math_module),\n",
    "    'interpolate':(interpolate_path, math_module),\n",
    "    'extrapolate':(extrapolate_path, math_module)\n",
    "           }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of set interpolate is 90000\n",
      "Length of set extrapolate is 60000\n",
      "Length of set train is 5999994\n",
      "CPU times: user 2.57 s, sys: 576 ms, total: 3.14 s\n",
      "Wall time: 3.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "input_texts = {}\n",
    "target_texts = {}\n",
    "\n",
    "for k, v in datasets.items():\n",
    "    input_texts[k], target_texts[k] = concatenate_texts(v[0], v[1])\n",
    "    print('Length of set {} is {}'.format(k, len(input_texts[k])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Sample:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INPUT: What is 2 - (1 + -5) - 11?\n",
      "OUTPUT: -5\n"
     ]
    }
   ],
   "source": [
    "print('INPUT:', input_texts['train'][42])\n",
    "print('OUTPUT:', target_texts['train'][42].strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Concatenate texts to get text metrics (max length, number of unique tokens, etc.):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_input_texts = sum(input_texts.values(), [])\n",
    "all_target_texts = sum(target_texts.values(), [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_characters = set(''.join(all_input_texts))\n",
    "target_characters = set(''.join(all_target_texts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of samples: 6149994\n",
      "Number of unique input tokens: 54\n",
      "Number of unique output tokens: 30\n",
      "Max sequence length for inputs: 160\n",
      "Max sequence length for outputs: 32\n"
     ]
    }
   ],
   "source": [
    "input_characters = sorted(list(input_characters))\n",
    "target_characters = sorted(list(target_characters))\n",
    "num_encoder_tokens = len(input_characters)\n",
    "num_decoder_tokens = len(target_characters)\n",
    "max_encoder_seq_length = max([len(txt) for txt in all_input_texts])\n",
    "max_decoder_seq_length = max([len(txt) for txt in all_target_texts])\n",
    "\n",
    "print('Number of samples:', len(all_input_texts))\n",
    "print('Number of unique input tokens:', num_encoder_tokens)\n",
    "print('Number of unique output tokens:', num_decoder_tokens)\n",
    "print('Max sequence length for inputs:', max_encoder_seq_length)\n",
    "print('Max sequence length for outputs:', max_decoder_seq_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete all texts to realease memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "del all_input_texts\n",
    "del all_target_texts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create train test splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_texts_train, input_texts_valid, target_texts_train, target_texts_valid = train_test_split(input_texts['train'], target_texts['train'], test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of training samples: 4799995\n"
     ]
    }
   ],
   "source": [
    "print('Number of training samples:', len(input_texts_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of validation samples: 1199999\n"
     ]
    }
   ],
   "source": [
    "print('Number of validation samples:', len(input_texts_valid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorise the text\n",
    "Before training, we need to map strings to a numerical representation. Create two lookup tables: one mapping question characters to numbers, and another for answer characters to number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a mapping from unique characters to indices\n",
    "input_token_index = dict([(char, i) for i, char in enumerate(input_characters)])\n",
    "target_token_index = dict([(char, i) for i, char in enumerate(target_characters)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'\\t': 0,\n",
       " '\\n': 1,\n",
       " ' ': 2,\n",
       " '(': 3,\n",
       " ')': 4,\n",
       " '*': 5,\n",
       " '+': 6,\n",
       " '-': 7,\n",
       " '.': 8,\n",
       " '/': 9,\n",
       " '0': 10,\n",
       " '1': 11,\n",
       " '2': 12,\n",
       " '3': 13,\n",
       " '4': 14,\n",
       " '5': 15,\n",
       " '6': 16,\n",
       " '7': 17,\n",
       " '8': 18,\n",
       " '9': 19,\n",
       " 'a': 20,\n",
       " 'b': 21,\n",
       " 'c': 22,\n",
       " 'd': 23,\n",
       " 'e': 24,\n",
       " 'f': 25,\n",
       " 'q': 26,\n",
       " 'r': 27,\n",
       " 's': 28,\n",
       " 't': 29}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_token_index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create keras data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "params = {'batch_size': settings_dict[\"batch_size\"],\n",
    "          'max_encoder_seq_length': max_encoder_seq_length,\n",
    "          'max_decoder_seq_length': max_decoder_seq_length,\n",
    "          'num_encoder_tokens': num_encoder_tokens,\n",
    "          'num_decoder_tokens': num_decoder_tokens,\n",
    "          'input_token_index': input_token_index,\n",
    "          'target_token_index': target_token_index,\n",
    "          'num_thinking_steps': settings_dict[\"thinking_steps\"]\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_generator = DataGenerator(input_texts=input_texts_train, target_texts=target_texts_train, **params)\n",
    "validation_generator = DataGenerator(input_texts=input_texts_valid, target_texts=target_texts_valid, **params)\n",
    "interpolate_generator = DataGenerator(input_texts=input_texts['interpolate'], target_texts=target_texts['interpolate'], **params)\n",
    "extrapolate_generator = DataGenerator(input_texts=input_texts['extrapolate'], target_texts=target_texts['extrapolate'], **params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_dict = {\n",
    "    'validation':validation_generator,\n",
    "    'interpolation': interpolate_generator,\n",
    "    'extrapolation': extrapolate_generator\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = NValidationSetsCallback(valid_dict)\n",
    "gradient = GradientLogger(live_metrics=['loss', 'exact_match_metric'], live_gaps=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = settings_dict['epochs']  # Number of epochs to train for.\n",
    "latent_dim = settings_dict['latent_dim']  # Latent dimensionality of the encoding space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm = LSTM_S2S(num_encoder_tokens, num_decoder_tokens, latent_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lstm.get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4687/4687 [==============================] - 826s 176ms/step - loss: 0.2631 - exact_match_metric: 0.0147\n"
     ]
    }
   ],
   "source": [
    "adam = Adam(lr=6e-4, beta_1=0.9, beta_2=0.995, epsilon=1e-9, decay=0.0, amsgrad=False, clipnorm=0.1)\n",
    "\n",
    "model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=[exact_match_metric])\n",
    "print('start training...')\n",
    "train_hist = model.fit_generator(training_generator,\n",
    "                                 epochs=epochs,\n",
    "                                 #use_multiprocessing=True, workers=8,\n",
    "                                 callbacks=[history, gradient],\n",
    "                                 verbose=0,\n",
    "                                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-1-aabe21b7fb28>, line 13)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-1-aabe21b7fb28>\"\u001b[0;36m, line \u001b[0;32m13\u001b[0m\n\u001b[0;31m    +)\u001b[0m\n\u001b[0m     ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "plt.plot(train_hist.history['loss'],color='C0', label='train')\n",
    "plt.plot(train_hist.history['validation_loss'], color='C0', label='valid', linestyle='--')\n",
    "plt.plot(train_hist.history['extrapolation_loss'], color='C1', label='extra',)\n",
    "plt.plot(train_hist.history['interpolation_loss'], color='C2', label='inter')\n",
    "\n",
    "plt.xlabel('epochs')\n",
    "plt.ylabel('loss')\n",
    "plt.legend(loc='best')\n",
    "plt.ylim([0,1])\n",
    "plt.grid(True, linestyle='--')\n",
    "plt.tight_layout()\n",
    "plt.savefig(settings_dict['save_path'] + 'losses.png', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(train_hist.history['exact_match_metric'],color='C0', label='train')\n",
    "plt.plot(train_hist.history['validation_exact_match_metric'], color='C0', label='valid', linestyle='--')\n",
    "plt.plot(train_hist.history['extrapolation_exact_match_metric'], color='C1', label='extra',)\n",
    "plt.plot(train_hist.history['interpolation_exact_match_metric'], color='C2', label='inter')\n",
    "\n",
    "plt.xlabel('epochs')\n",
    "plt.ylabel('exact match metric')\n",
    "plt.legend(loc='best')\n",
    "plt.ylim([0,1])\n",
    "plt.grid(True, linestyle='--')\n",
    "plt.tight_layout()\n",
    "plt.savefig(settings_dict['save_path'] + 'metrics.png', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(settings_dict['save_path']+'experiments_output.pkl','wb') as file:\n",
    "    pickle.dump(train_hist.history, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(settings_dict['save_path']+'this_model.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: '../../artifacts/settings.json'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-78-073251a7b242>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'../..'\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0msettings_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'save_path'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'settings.json'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'wb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfile\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m     \u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdump\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msettings_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: '../../artifacts/settings.json'"
     ]
    }
   ],
   "source": [
    "with open(settings_dict['save_path']+'settings.json','w') as file:\n",
    "    json.dump(settings_dict, file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
